local component = require("component")
local computer = require("computer")
local event = require("event")
local thread = require("thread")
local serialization = require("serialization")
local filesystem = require("filesystem")

-- ========== 配置区域 ==========
local BIND_PORT = 67       -- DHCP 服务端口
local CLIENT_PORT = 68     -- DHCP 客户端端口
local ARP_PORT = 69        -- ARP协议端口
local ICMP_PORT = 70       -- ICMP协议端口
local LEASE_TIME = 3600    -- IP 租约时间(秒)
local IP_POOL = {          -- IP 地址池
    "192.168.1.100", "192.168.1.101", "192.168.1.102",
    "192.168.1.103", "192.168.1.104", "192.168.1.105"
}
local SERVER_IP = "192.168.1.1"  -- 服务器/网关IP
local CONFIG_FILE = "/etc/dhcp_server.cfg"
local ARP_CACHE_TIME = 300 -- ARP缓存5分钟
-- =============================

local modem = nil
local leases = {}      -- IP 租约表: {[MAC] = {ip, expire_time, sender}}
local arpTable = {}    -- ARP表: {[IP] = {mac, timestamp}}
local running = true
local dhcpThread = nil
local arpThread = nil
local icmpThread = nil

local serverStats = {
    dhcp_requests = 0,
    dhcp_offers = 0,
    dhcp_acks = 0,
    arp_requests = 0,
    arp_replies = 0,
    icmp_requests = 0,
    icmp_replies = 0,
    start_time = 0
}

-- ===============================
-- 工具函数
-- ===============================

local function saveServerConfig()
    local config = {
        leases = leases,
        arpTable = arpTable,
        stats = serverStats,
        server_ip = SERVER_IP
    }

    local dir = filesystem.path(CONFIG_FILE)
    if not filesystem.exists(dir) then
        filesystem.makeDirectory(dir)
    end

    local file = io.open(CONFIG_FILE, "w")
    if file then
        file:write(serialization.serialize(config))
        file:close()
        return true
    end
    return false
end

local function loadServerConfig()
    if filesystem.exists(CONFIG_FILE) then
        local file = io.open(CONFIG_FILE, "r")
        if file then
            local content = file:read("*a")
            file:close()
            local success, config = pcall(serialization.unserialize, content)
            if success and config then
                leases = config.leases or {}
                arpTable = config.arpTable or {}
                serverStats = config.stats or serverStats
                return true
            end
        end
    end
    return false
end

-- 网卡选择功能
local function selectNetworkCard()
    local modems = {}
    for address in component.list("modem") do
        table.insert(modems, {address = address, proxy = component.proxy(address)})
    end

    if #modems == 0 then
        print("Error: No network card found!")
        return nil
    elseif #modems == 1 then
        print("Using network card: " .. modems[1].address:sub(1,8) .. "...")
        return modems[1].proxy
    else
        print("Multiple network cards detected:")
        for i, modem in ipairs(modems) do
            local isWireless = modem.isWireless and modem.isWireless() or false
            local cardType = isWireless and "(Wireless)" or "(Wired)"
            print(string.format("%d. %s %s", i, modem.address, cardType))
        end

        io.write("Enter card number for DHCP server (1-" .. #modems .. "): ")
        local choice = tonumber(io.read())

        if choice and choice >= 1 and choice <= #modems then
            print("Selected: " .. modems[choice].address:sub(1,8) .. "...")
            return modems[choice].proxy
        else
            print("Invalid selection, using first card")
            return modems[1].proxy
        end
    end
end

-- 清理过期租约
local function cleanupExpiredLeases()
    local now = computer.uptime()
    local cleaned = 0
    for mac, lease in pairs(leases) do
        if mac ~= "server" and lease.expire < now then
            print("[DHCP] Lease expired for " .. mac:sub(1,8) .. "... (IP: " .. lease.ip .. ")")

            -- 同时清理ARP表中对应的条目
            arpTable[lease.ip] = nil
            leases[mac] = nil
            cleaned = cleaned + 1
        end
    end
    return cleaned
end

-- 清理ARP表
local function cleanupArpTable()
    local now = computer.uptime()
    for ip, entry in pairs(arpTable) do
        if now - entry.timestamp > ARP_CACHE_TIME then
            arpTable[ip] = nil
        end
    end
end

-- IP地址工具函数
local function isInSameSubnet(ip1, ip2, netmask)
    -- 简化实现：只支持/24网络
    local subnet1 = ip1:match("(%d+%.%d+%.%d+%.)%d+")
    local subnet2 = ip2:match("(%d+%.%d+%.%d+%.)%d+")
    return subnet1 == subnet2
end

-- 根据IP查找MAC地址
local function findMACByIP(targetIP)
    cleanupArpTable()

    -- 首先检查ARP表
    if arpTable[targetIP] then
        return arpTable[targetIP].mac
    end

    -- 检查DHCP租约表
    for mac, lease in pairs(leases) do
        if lease.ip == targetIP and lease.expire > computer.uptime() then
            return mac
        end
    end

    return nil
end

-- ===============================
-- DHCP 服务器功能
-- ===============================

local function dhcpListener()
    while running do
        local eventData = {event.pull(0.5, "modem_message")}

        if eventData[1] == "modem_message" then
            local _, _, sender, port, _, payload = table.unpack(eventData)

            if port == BIND_PORT and payload then
                local success, msg = pcall(serialization.unserialize, payload)

                if success and msg and msg.type and msg.mac then
                    local mac = msg.mac
                    local client_port = msg.client_port or CLIENT_PORT

                    serverStats.dhcp_requests = serverStats.dhcp_requests + 1
                    print("[DHCP] Received " .. msg.type .. " from " .. mac:sub(1,8) .. "...")

                    -- DHCP DISCOVER 处理
                    if msg.type == "DISCOVER" then
                        cleanupExpiredLeases()

                        local assignedIp = nil

                        -- 检查是否已有租约
                        if leases[mac] and leases[mac].expire > computer.uptime() then
                            assignedIp = leases[mac].ip
                            print("[DHCP] Renewing existing lease: " .. assignedIp)
                        else
                            -- 查找新的IP
                            for _, ip in ipairs(IP_POOL) do
                                local alreadyAssigned = false
                                for leaseMac, lease in pairs(leases) do
                                    if leaseMac ~= "server" and lease.ip == ip and lease.expire > computer.uptime() then
                                        alreadyAssigned = true
                                        break
                                    end
                                end

                                if not alreadyAssigned then
                                    assignedIp = ip
                                    break
                                end
                            end
                        end

                        -- 发送 OFFER
                        if assignedIp then
                            leases[mac] = {
                                ip = assignedIp,
                                expire = computer.uptime() + LEASE_TIME,
                                sender = sender
                            }

                            -- 更新ARP表
                            arpTable[assignedIp] = {
                                mac = mac,
                                timestamp = computer.uptime()
                            }

                            local offer = {
                                type = "OFFER",
                                mac = mac,
                                ip = assignedIp,
                                netmask = "255.255.255.0",
                                gateway = SERVER_IP,
                                dns = SERVER_IP,
                                lease_time = LEASE_TIME
                            }

                            modem.send(sender, client_port, serialization.serialize(offer))
                            serverStats.dhcp_offers = serverStats.dhcp_offers + 1
                            print("[DHCP] Offered IP " .. assignedIp .. " to " .. mac:sub(1,8) .. "...")
                        else
                            print("[DHCP] No available IP for " .. mac:sub(1,8) .. "...")
                        end

                        -- DHCP REQUEST 处理
                    elseif msg.type == "REQUEST" then
                        if leases[mac] and leases[mac].ip == msg.requested_ip then
                            leases[mac].expire = computer.uptime() + LEASE_TIME
                            leases[mac].sender = sender

                            -- 更新ARP表
                            arpTable[msg.requested_ip] = {
                                mac = mac,
                                timestamp = computer.uptime()
                            }

                            modem.send(sender, client_port, serialization.serialize({
                                type = "ACK",
                                mac = mac,
                                ip = leases[mac].ip
                            }))
                            serverStats.dhcp_acks = serverStats.dhcp_acks + 1
                            print("[DHCP] Confirmed IP " .. leases[mac].ip .. " for " .. mac:sub(1,8) .. "...")
                        else
                            print("[DHCP] REQUEST rejected for " .. mac:sub(1,8) .. "... (IP: " .. (msg.requested_ip or "unknown") .. ")")
                        end
                    end
                else
                    print("[DHCP] Error parsing packet from " .. sender)
                end
            end
        end

        -- 定期保存配置和清理
        if math.floor(computer.uptime()) % 300 == 0 then -- 每5分钟
            cleanupExpiredLeases()
            cleanupArpTable()
            saveServerConfig()
        end
    end
end

-- ===============================
-- ARP 协议实现
-- ===============================

local function arpListener()
    while running do
        local eventData = {event.pull(0.5, "modem_message")}

        if eventData[1] == "modem_message" then
            local _, _, sender, port, _, payload = table.unpack(eventData)

            if port == ARP_PORT and payload then
                local success, packet = pcall(serialization.unserialize, payload)

                if success and packet then
                    -- 处理ARP请求
                    if packet.type == "arp_request" then
                        serverStats.arp_requests = serverStats.arp_requests + 1

                        -- 更新发送者信息到ARP表
                        arpTable[packet.sender_ip] = {
                            mac = packet.sender_mac,
                            timestamp = computer.uptime()
                        }

                        print("[ARP] ARP request from " .. packet.sender_ip ..
                                " (" .. packet.sender_mac:sub(1,8) .. "...) looking for " .. packet.target_ip)

                        -- 如果目标IP是服务器自己
                        if packet.target_ip == SERVER_IP then
                            modem.send(packet.sender_mac, ARP_PORT, serialization.serialize({
                                type = "arp_reply",
                                sender_ip = SERVER_IP,
                                sender_mac = modem.address,
                                target_ip = packet.sender_ip,
                                target_mac = packet.sender_mac
                            }))
                            serverStats.arp_replies = serverStats.arp_replies + 1
                            print("[ARP] Replied: " .. SERVER_IP .. " is at " .. modem.address:sub(1,8) .. "...")

                            -- 如果目标IP在DHCP租约表中，转发请求
                        else
                            local targetMAC = findMACByIP(packet.target_ip)
                            if targetMAC then
                                print("[ARP] Forwarding ARP request to " .. packet.target_ip)
                                modem.send(targetMAC, ARP_PORT, serialization.serialize(packet))
                            else
                                print("[ARP] No route to " .. packet.target_ip)
                            end
                        end

                        -- 处理ARP回复，转发给请求者
                    elseif packet.type == "arp_reply" then
                        -- 更新ARP表
                        arpTable[packet.sender_ip] = {
                            mac = packet.sender_mac,
                            timestamp = computer.uptime()
                        }

                        -- 如果回复的目标不是服务器，转发回复
                        if packet.target_ip ~= SERVER_IP then
                            local targetMAC = findMACByIP(packet.target_ip)
                            if targetMAC then
                                print("[ARP] Forwarding ARP reply from " .. packet.sender_ip .. " to " .. packet.target_ip)
                                modem.send(targetMAC, ARP_PORT, serialization.serialize(packet))
                            end
                        end
                    end
                end
            end
        end
    end
end

-- ===============================
-- ICMP 协议实现
-- ===============================

local function icmpListener()
    while running do
        local eventData = {event.pull(0.5, "modem_message")}

        if eventData[1] == "modem_message" then
            local _, _, sender, port, _, payload = table.unpack(eventData)

            if port == ICMP_PORT and payload then
                local success, packet = pcall(serialization.unserialize, payload)

                if success and packet then
                    -- 处理ICMP Echo请求
                    if packet.type == "icmp_echo_request" then
                        serverStats.icmp_requests = serverStats.icmp_requests + 1

                        print("[ICMP] Ping from " .. packet.source_ip .. " to " .. packet.target_ip ..
                                " (seq=" .. (packet.seq or "?") .. ")")

                        -- 如果目标是服务器自己
                        if packet.target_ip == SERVER_IP then
                            modem.send(packet.source_mac, ICMP_PORT, serialization.serialize({
                                type = "icmp_echo_reply",
                                target_ip = packet.source_ip,
                                source_ip = SERVER_IP,
                                source_mac = modem.address,
                                target_mac = packet.source_mac,
                                id = packet.id,
                                seq = packet.seq,
                                timestamp = packet.timestamp
                            }))
                            serverStats.icmp_replies = serverStats.icmp_replies + 1
                            print("[ICMP] Replied to ping from " .. packet.source_ip)

                            -- 如果目标是其他客户端，转发请求
                        else
                            local targetMAC = findMACByIP(packet.target_ip)
                            if targetMAC then
                                print("[ICMP] Forwarding ping from " .. packet.source_ip .. " to " .. packet.target_ip)
                                modem.send(targetMAC, ICMP_PORT, serialization.serialize(packet))
                            else
                                print("[ICMP] No route to " .. packet.target_ip)

                                -- 发送目的地不可达
                                modem.send(packet.source_mac, ICMP_PORT, serialization.serialize({
                                    type = "icmp_dest_unreachable",
                                    target_ip = packet.source_ip,
                                    source_ip = SERVER_IP,
                                    source_mac = modem.address,
                                    target_mac = packet.source_mac,
                                    original_target = packet.target_ip,
                                    id = packet.id,
                                    seq = packet.seq
                                }))
                            end
                        end

                        -- 处理ICMP Echo回复，转发给请求者
                    elseif packet.type == "icmp_echo_reply" then
                        if packet.target_ip ~= SERVER_IP then
                            local targetMAC = findMACByIP(packet.target_ip)
                            if targetMAC then
                                print("[ICMP] Forwarding ping reply from " .. packet.source_ip .. " to " .. packet.target_ip)
                                modem.send(targetMAC, ICMP_PORT, serialization.serialize(packet))
                            end
                        end
                    end
                end
            end
        end
    end
end

-- ===============================
-- 服务器端Ping功能
-- ===============================

local function arpRequest(targetIP)
    cleanupArpTable()

    -- 检查ARP缓存
    if arpTable[targetIP] and (computer.uptime() - arpTable[targetIP].timestamp < ARP_CACHE_TIME) then
        return arpTable[targetIP].mac
    end

    print("[ARP] Resolving " .. targetIP .. "...")

    -- 发送ARP请求
    modem.broadcast(ARP_PORT, serialization.serialize({
        type = "arp_request",
        sender_ip = SERVER_IP,
        sender_mac = modem.address,
        target_ip = targetIP
    }))

    -- 等待ARP回复
    local deadline = computer.uptime() + 2
    while computer.uptime() < deadline do
        local eventData = {event.pull(0.1, "modem_message")}
        if eventData[1] == "modem_message" then
            local _, _, from, port, _, payload = table.unpack(eventData)

            if port == ARP_PORT and payload then
                local success, packet = pcall(serialization.unserialize, payload)
                if success and packet and packet.type == "arp_reply" and
                        packet.target_ip == SERVER_IP and packet.sender_ip == targetIP then

                    arpTable[targetIP] = {
                        mac = packet.sender_mac,
                        timestamp = computer.uptime()
                    }

                    print("[ARP] " .. targetIP .. " is at " .. packet.sender_mac:sub(1,8) .. "...")
                    return packet.sender_mac
                end
            end
        end
    end

    print("[ARP] No response from " .. targetIP)
    return nil
end

local function pingFromServer(targetIP, count)
    count = count or 4
    print("PING " .. targetIP .. " from server " .. SERVER_IP)

    -- 进行ARP解析
    local targetMAC = arpRequest(targetIP)
    if not targetMAC then
        print("ARP failed - cannot reach " .. targetIP)
        return
    end

    local sent = 0
    local received = 0
    local responses = {}

    for i = 1, count do
        local startTime = computer.uptime()
        local pingId = math.random(1, 65535)

        modem.send(targetMAC, ICMP_PORT, serialization.serialize({
            type = "icmp_echo_request",
            target_ip = targetIP,
            source_ip = SERVER_IP,
            source_mac = modem.address,
            id = pingId,
            seq = i,
            timestamp = startTime
        }))

        sent = sent + 1
        local replied = false

        local deadline = startTime + 3
        while computer.uptime() < deadline and not replied do
            local eventData = {event.pull(0.1, "modem_message")}
            if eventData[1] == "modem_message" then
                local _, _, from, port, _, payload = table.unpack(eventData)
                if port == ICMP_PORT and payload then
                    local success, packet = pcall(serialization.unserialize, payload)
                    if success and packet and packet.id == pingId and
                            packet.source_ip == targetIP and packet.target_ip == SERVER_IP then

                        if packet.type == "icmp_echo_reply" then
                            local responseTime = (computer.uptime() - startTime) * 1000
                            table.insert(responses, responseTime)
                            print(string.format("Reply from %s: seq=%d time=%.1fms",
                                    targetIP, packet.seq, responseTime))
                            received = received + 1
                            replied = true

                        elseif packet.type == "icmp_dest_unreachable" then
                            print("Destination unreachable for seq=" .. i)
                            replied = true
                        end
                    end
                end
            end
        end

        if not replied then
            print("Request timeout for seq=" .. i)
        end

        if i < count then
            os.sleep(1)
        end
    end

    print(string.format("\n%d packets transmitted, %d received, %.1f%% packet loss",
            sent, received, (sent - received) / sent * 100))

    if #responses > 0 then
        local total = 0
        local min_time = responses[1]
        local max_time = responses[1]

        for _, time in ipairs(responses) do
            total = total + time
            if time < min_time then min_time = time end
            if time > max_time then max_time = time end
        end

        local avg_time = total / #responses
        print(string.format("round-trip min/avg/max = %.1f/%.1f/%.1f ms",
                min_time, avg_time, max_time))
    end
end

-- ===============================
-- 统计和监控功能
-- ===============================

local function showStats()
    local uptime = computer.uptime() - serverStats.start_time
    local hours = math.floor(uptime / 3600)
    local minutes = math.floor((uptime % 3600) / 60)
    local seconds = math.floor(uptime % 60)

    print("\n===== Server Statistics =====")
    print("Server IP    : " .. SERVER_IP)
    print("Network Card : " .. modem.address)
    print("Uptime       : " .. string.format("%02d:%02d:%02d", hours, minutes, seconds))
    print("DHCP Requests: " .. serverStats.dhcp_requests)
    print("DHCP Offers  : " .. serverStats.dhcp_offers)
    print("DHCP ACKs    : " .. serverStats.dhcp_acks)
    print("ARP Requests : " .. serverStats.arp_requests)
    print("ARP Replies  : " .. serverStats.arp_replies)
    print("ICMP Requests: " .. serverStats.icmp_requests)
    print("ICMP Replies : " .. serverStats.icmp_replies)
    print("=============================\n")
end

local function showLeases()
    cleanupExpiredLeases()
    print("\n=== Active DHCP Leases ===")
    print("MAC Address     \tIP Address    \tExpires In\tLast Seen")
    print("------------------------------------------------------------")

    local now = computer.uptime()
    local count = 0
    for mac, lease in pairs(leases) do
        if mac ~= "server" then
            local remaining = math.max(0, lease.expire - now)
            local lastSeen = lease.sender or "Unknown"
            print(string.format("%-16s\t%-14s\t%d mins\t%s",
                    mac:sub(1,12) .. "...",
                    lease.ip,
                    math.floor(remaining / 60),
                    lastSeen:sub(1,8) .. "..."))
            count = count + 1
        end
    end

    if count == 0 then
        print("No active leases")
    end
    print("============================================\n")
end

local function showArpTable()
    cleanupArpTable()
    print("\n=== ARP Table ===")
    print("IP Address      MAC Address       Age")
    print("--------------------------------------")

    local now = computer.uptime()
    local count = 0
    for ip, entry in pairs(arpTable) do
        local age = math.floor(now - entry.timestamp)
        print(string.format("%-15s %-17s %d sec", ip, entry.mac:sub(1,17) .. "...", age))
        count = count + 1
    end

    if count == 0 then
        print("No entries")
    end
    print("==========================\n")
end

-- ===============================
-- CLI 命令处理
-- ===============================

local function showHelp()
    print("Available commands:")
    print("  list      - Show active DHCP leases")
    print("  arp       - Show ARP table")
    print("  stats     - Show server statistics")
    print("  ping <ip> [count] - Ping a client")
    print("  route     - Show routing table")
    print("  clean     - Clean expired leases and ARP entries")
    print("  save      - Save configuration")
    print("  reload    - Reload configuration")
    print("  kick <ip> - Kick client (revoke lease)")
    print("  reset     - Reset all statistics")
    print("  exit      - Stop server")
    print("  help      - Show this help")
end

local function showRoutes()
    print("\n=== Routing Table ===")
    print("Destination     Gateway         Interface")
    print("----------------------------------------")
    print("192.168.1.0/24  0.0.0.0         " .. modem.address:sub(1,8) .. "... (direct)")
    print("0.0.0.0/0       " .. SERVER_IP .. "         " .. modem.address:sub(1,8) .. "... (default)")
    print("=======================\n")
end

local function handleCommand(input)
    local parts = {}
    for word in input:gmatch("%S+") do
        table.insert(parts, word:lower())
    end

    if #parts == 0 then
        return true
    end

    local cmd = parts[1]

    if cmd == "help" then
        showHelp()

    elseif cmd == "list" then
        showLeases()

    elseif cmd == "arp" then
        showArpTable()

    elseif cmd == "stats" then
        showStats()

    elseif cmd == "ping" then
        if parts[2] then
            local count = tonumber(parts[3]) or 4
            pingFromServer(parts[2], count)
        else
            print("Usage: ping <ip> [count]")
        end

    elseif cmd == "route" then
        showRoutes()

    elseif cmd == "clean" then
        local leases_cleaned = cleanupExpiredLeases()
        cleanupArpTable()
        print("Cleaned " .. leases_cleaned .. " expired leases and old ARP entries")

    elseif cmd == "save" then
        if saveServerConfig() then
            print("Configuration saved")
        else
            print("Failed to save configuration")
        end

    elseif cmd == "reload" then
        if loadServerConfig() then
            print("Configuration reloaded")
        else
            print("Failed to load configuration")
        end

    elseif cmd == "kick" then
        if parts[2] then
            local targetIP = parts[2]
            for mac, lease in pairs(leases) do
                if lease.ip == targetIP then
                    leases[mac] = nil
                    arpTable[targetIP] = nil
                    print("Kicked client with IP: " .. targetIP)
                    return true
                end
            end
            print("Client with IP " .. targetIP .. " not found")
        else
            print("Usage: kick <ip>")
        end

    elseif cmd == "reset" then
        serverStats = {
            dhcp_requests = 0,
            dhcp_offers = 0,
            dhcp_acks = 0,
            arp_requests = 0,
            arp_replies = 0,
            icmp_requests = 0,
            icmp_replies = 0,
            start_time = computer.uptime()
        }
        print("Statistics reset")

    elseif cmd == "exit" then
        return false

    else
        print("Unknown command: " .. cmd .. " (type 'help' for available commands)")
    end

    return true
end

-- ===============================
-- 主程序
-- ===============================

local function main()
    -- 初始化
    modem = selectNetworkCard()
    if not modem then
        print("Failed to initialize network card")
        return
    end

    -- 打开所有协议端口
    modem.open(BIND_PORT)  -- DHCP端口
    modem.open(ARP_PORT)   -- ARP端口
    modem.open(ICMP_PORT)  -- ICMP端口

    -- 加载配置
    loadServerConfig()

    -- 设置服务器自己的IP和ARP条目
    leases["server"] = {ip = SERVER_IP, expire = math.huge}
    arpTable[SERVER_IP] = {mac = modem.address, timestamp = computer.uptime()}

    -- 初始化统计
    serverStats.start_time = computer.uptime()

    print("DHCP/Network Server Started")
    print("Server IP: " .. SERVER_IP)
    print("Network card: " .. modem.address)
    print("Protocols: DHCP(" .. BIND_PORT .. "), ARP(" .. ARP_PORT .. "), ICMP(" .. ICMP_PORT .. ")")
    print("Type 'help' for available commands\n")

    -- 启动后台协议处理线程
    dhcpThread = thread.create(dhcpListener)
    arpThread = thread.create(arpListener)
    icmpThread = thread.create(icmpListener)

    -- CLI主循环
    while running do
        io.write("netserver> ")
        local input = io.read()

        if not input then
            break
        end

        if not handleCommand(input) then
            break
        end
    end

    -- 清理
    running = false
    if dhcpThread then dhcpThread:kill() end
    if arpThread then arpThread:kill() end
    if icmpThread then icmpThread:kill() end
    saveServerConfig()
    modem.close(BIND_PORT)
    modem.close(ARP_PORT)
    modem.close(ICMP_PORT)
    print("Server stopped")
end

-- 启动主程序
main()
